# # -*- coding: utf-8 -*-
# """
# Created on Tue Dec 29 14:15:43 2020

# @author: 26297
# """

# import torch
# import torch.nn as nn
# import torch.optim as optim
# from torchvision import datasets, transforms, models
# from torch.utils.data import DataLoader
# from age_net import AgeNet
# import matplotlib.pyplot as plt
# import torch.nn.functional as F
# from time import time
# import os
# transform = transforms.Compose([
#             transforms.Resize(112),
#             transforms.ToTensor()
#             ])
# root = '../trainSet'
# train_dataset = datasets.ImageFolder(root, transform)
# # sample_batch = [1,2,3,4,5,6,7,8,9]
# train_loader = torch.utils.data.DataLoader(train_dataset,num_workers=8, batch_size=8, shuffle=True)
# # train_loader = torch.utils.data.DataLoader(train_dataset,num_workers=8, batch_size=8)

# # import torch
# # import 

# # class MySample(Sampler):
# #     def __init__(self,data_source,indices):
# #         super(MySampler,self).__init__(data_source, )
# #         self.indices = indices
# from torch.utils.data.sampler import Sampler
# class CustomSampler(Sampler):
#     def __init__(self, data):
#         self.data = data

#     def __iter__(self):
#         indices = []
#         for n in range(self.data.num_classes):
#             index = torch.where(self.data.label == n)[0]
#             indices.append(index)
#         indices = torch.cat(indices, dim=0)
#         return iter(indices)

#     def __len__(self):
#         return len(self.data)
    
# sampler = CustomSampler()
# transform = transforms.Compose([
#             transforms.Resize(112),
#             transforms.ToTensor()
#             ])
# root = '../trainSet'
# train_dataset = datasets.ImageFolder(root, transform)
# # sample_batch = [1,2,3,4,5,6,7,8,9]
# train_loader = torch.utils.data.DataLoader(train_dataset,sampler=CustomSampler)
import numpy as np
from torch.utils.data.sampler import WeightedRandomSampler
samples_count = [1647,2845,1512,4900,4900,4900,4900,2866,1382,412]
samples_weight = 1/(samples_count/np.sum(samples_count))

samples_set = [samples_weight[k] for k in range(0,len(samples_count)) for i in range(0,samples_count[k])]
samples_num = 10

sampler = WeightedRandomSampler(samples_set,samples_num)




